import matplotlib.pyplot as plt
import matplotlib.patches as patches

# 创建一个新的图形
fig, ax = plt.subplots(figsize=(12, 8))

# 定义每个阶段的层的名称、通道数、层数和步长
stages = [
    {"name": "Conv3x3", "channels": 24, "layers": 1, "stride": 2},
    {"name": "Fused-MBConv1, k3x3", "channels": 24, "layers": 2, "stride": 1},
    {"name": "Fused-MBConv4, k3x3", "channels": 48, "layers": 4, "stride": 2},
    {"name": "Fused-MBConv4, k3x3", "channels": 64, "layers": 4, "stride": 2},
    {"name": "MBConv4, k3x3, SE0.25", "channels": 128, "layers": 6, "stride": 2},
    {"name": "MBConv6, k3x3, SE0.25", "channels": 160, "layers": 9, "stride": 1},
    {"name": "MBConv6, k3x3, SE0.25", "channels": 272, "layers": 15, "stride": 2},
    {"name": "Conv1x1 & Pooling & FC", "channels": 1792, "layers": 1, "stride": None}
]

# 计算每个阶段的宽度
stage_widths = [stage["layers"] for stage in stages]
total_width = sum(stage_widths)

# 设置每个阶段的x位置
x_positions = []
current_x = 0
for width in stage_widths:
    x_positions.append(current_x + width / 2)
    current_x += width

# 绘制每个阶段的矩形
y_position = 0
for i, stage in enumerate(stages):
    rect = patches.Rectangle((x_positions[i] - stage["layers"] / 2, y_position), stage["layers"], 1, linewidth=1, edgecolor='black', facecolor='lightblue')
    ax.add_patch(rect)
    plt.text(x_positions[i], y_position + 0.5, f"{stage['name']}\n#Channels {stage['channels']}", ha='center', va='center', fontsize=8)

    # 更新y位置以绘制下一个阶段
    y_position += 1

# 设置图形的显示范围
ax.set_xlim(-0.5, total_width - 0.5)
ax.set_ylim(-1, len(stages))

# 添加箭头表示数据流
for i in range(len(stages) - 1):
    ax.annotate('', xy=(x_positions[i], 0.5), xytext=(x_positions[i + 1], 0.5),
                arrowprops=dict(facecolor='black', shrink=0.05))

# 显示图形
plt.title("Network Structure")
plt.axis('off')  # 不显示坐标轴
plt.show()